import numpy as np
from scipy.integrate import odeint
from scipy.constants import physical_constants

WIDTHS = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 12, 14, 16, 18, 20, 25, 30, 35, 40, 45, 50, 62.5, 75, 100]
RATIOS = [0, 10, 20, 30, 45, 60, 80, 100]
E_0 = (0.2 * 2.2 + 2.8 + 2 * 4.6) / 9.6 * 1e3
N_A = physical_constants['Avogadro constant'][0]
e = physical_constants['atomic unit of charge'][0]
m_e = physical_constants['electron mass'][0]
c = physical_constants['speed of light in vacuum'][0]

I_Si = 173e-6
Z_Si = 14
N_Si = 2.33 / 28.0855 * N_A

def Cal_S(N, Z, E, I):
    v = c * 1e2 * np.sqrt(1 - np.square(m_e / (m_e + E / np.square(c))))
    beta = v / c * 1e-2
    #B = Z * (np.log(2 * m_e * 1e-3 * np.square(v * 1e-2) * 1e-6 / e / I) - np.log(1 - np.square(beta)) - np.square(beta))
    something = 1e9 * I * e
    B = Z * (np.log(2 * m_e * np.square(v * 1e-2) / something) - np.log(1 - np.square(beta)) - np.square(beta))
    return (4 * np.pi * np.power(e, 4) * N * np.power(3e9, 4)) / (m_e * np.square(v) * 1.6e-6)

def Get_E(mat, E0, l):  #l单位cm，E0单位MeV
    if l == 0:
        return E0
    else:
        x_sequence = np.linspace(0, l, num=1000)
        if mat == 'Al':
            I = 166e-6
            Z = 13
            N = 2.699 / 26.981539 * N_A
        elif matt == 'Ta':
            I = 718e-6
            Z = 73
            N = 16.654 / 180.9479 * N_A

        def Cal_E(E, x):
            return - Cal_S(N, Z, E, I)
        
        sol = odeint(Cal_E, E0, x_sequence)
        return sol[:, 0][-1]

print(E_0)
print(Get_E('Al', E_0, 1))
D = []
for ratio in RATIOS:
    D.append([])
    for width in WIDTHS:
        if width == 0:
            S = Cal_S(N_Si, Z_Si, E_0, I_Si)
        else:
            E_Si = Get_E('Al', E_0, width * 0.9 * (1 - ratio / 100))
            E_Si = Get_E('Ta', E_Si, width * (ratio / 100))
            E_Si = Get_E('Al', E_Si, width * 0.1 * (1 - ratio / 100))
            S = Cal_S(N_Si, Z_Si, E_Si, I_Si)